from ultralytics import YOLO
import cv2
import os
import xml.etree.ElementTree as ET
from xml.dom import minidom

# 类别索引到名称的映射
CLASS_MAPPING = {
    0: 'hel',
    1: 'nohel',
    2: 'wheel'
}


def create_voc_xml(image_path, detections, width, height):
    """创建VOC格式的XML树"""
    root = ET.Element("annotation")

    # 文件名和路径信息
    folder = ET.SubElement(root, "folder")
    folder.text = os.path.basename(os.path.dirname(image_path))
    filename = ET.SubElement(root, "filename")
    filename.text = os.path.basename(image_path)
    path = ET.SubElement(root, "path")
    path.text = os.path.abspath(image_path)

    # 图像尺寸信息
    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = str(width)
    ET.SubElement(size, "height").text = str(height)
    ET.SubElement(size, "depth").text = "3"  # RGB图像

    # 添加检测对象
    for det in detections:
        obj = ET.SubElement(root, "object")
        # 使用检测到的类别名称（det[4]）
        ET.SubElement(obj, "name").text = det[4]
        ET.SubElement(obj, "pose").text = "Unspecified"
        ET.SubElement(obj, "truncated").text = "0"
        ET.SubElement(obj, "difficult").text = "0"

        # 添加边界框
        bbox = ET.SubElement(obj, "bndbox")
        ET.SubElement(bbox, "xmin").text = str(int(det[0]))
        ET.SubElement(bbox, "ymin").text = str(int(det[1]))
        ET.SubElement(bbox, "xmax").text = str(int(det[2]))
        ET.SubElement(bbox, "ymax").text = str(int(det[3]))

    # 格式化XML
    rough_xml = ET.tostring(root, 'utf-8')
    reparsed = minidom.parseString(rough_xml)
    return reparsed.toprettyxml(indent="  ")


def process_folder(folder_path, output_dir=None):
    """处理整个文件夹的图片"""
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']

    # 加载YOLOv8模型
    model = YOLO(r'D:\CodeCNN\yolov8-study\runs\detect\train27\weights\best.pt')

    # 遍历文件夹
    for file in os.listdir(folder_path):
        file_path = os.path.join(folder_path, file)

        # 检查是否为支持的图片文件
        if os.path.isfile(file_path) and any(file.lower().endswith(ext) for ext in image_extensions):
            print(f"Processing: {file_path}")

            # 读取图片获取尺寸
            img = cv2.imread(file_path)
            if img is None:
                print(f"  Failed to read image: {file_path}")
                continue

            height, width, _ = img.shape

            # 进行预测（检测所有类别）
            results = model.predict(source=file_path, conf=0.25)  # 移除了classes参数

            # 提取所有类别检测结果（0,1,2）
            detections = []
            for result in results:
                for box in result.boxes:
                    cls_idx = int(box.cls.item())  # 获取类别索引
                    if cls_idx in CLASS_MAPPING:  # 只处理目标类别
                        xyxy = box.xyxy[0].tolist()
                        # 添加类别名称到检测结果
                        class_name = CLASS_MAPPING[cls_idx]
                        detections.append([*xyxy, class_name])

            # 生成VOC XML
            xml_str = create_voc_xml(file_path, detections, width, height)

            # 设置输出路径
            if output_dir is None:
                output_dir = folder_path
            os.makedirs(output_dir, exist_ok=True)

            # 保存XML文件
            xml_filename = os.path.splitext(file)[0] + ".xml"
            xml_path = os.path.join(output_dir, xml_filename)

            with open(xml_path, 'w', encoding='utf-8') as f:
                f.write(xml_str)

            # 统计各类别检测数量
            count_stats = {name: sum(1 for d in detections if d[4] == name)
                           for name in CLASS_MAPPING.values()}

            print(f"  Saved VOC XML to: {xml_path}")
            print(f"  Detections: {count_stats}")

    print("\nBatch processing completed!")


if __name__ == "__main__":
    # 配置路径
    input_folder = r"D:\自动生成图片\0825\images"  # 替换为你的图片文件夹路径
    output_folder = r"D:\自动生成图片\0825\labels"  # XML输出目录

    # 处理整个文件夹
    process_folder(input_folder, output_folder)